
from graph import graph

graph.get_graph().print_ascii()
print(graph.get_graph().draw_mermaid())

def stream_graph_updates(user_input: str):
    for event in graph.stream({"messages": [("user", user_input)]}, stream_mode="values"):
        # event.pretty_print()
        # print('event: \n', event)
        for value in event.values():
            print('AI: \n', value[-1].content)
            # print("Assistant:", value["messages"][-1].content)

def do_stream(user_input):
    resps = graph.stream(
        {"messages": [("user", user_input)]}, stream_mode="values"
    )
    last = None
    for value in resps:
        # value["messages"][-1].pretty_print()
        # print("AI:", value["messages"][-1].content)
        last = value["messages"][-1].content
        pass

    print('AI: ', last)
    # print('AI: ', resps[-1]['messages'][-1].content)


while True:
    try:
        user_input = input("User: ")
        if user_input.lower() in ["quit", "exit", "q"]:
            print("Goodbye!")
            break
        
        do_stream(user_input)
        # stream_graph_updates(user_input)
    except Exception as e:
        print("Error: input() is not available: ", e)
        # fallback if input() is not available
        user_input = "What do you know about LangGraph?"
        print("User: " + user_input)
        stream_graph_updates(user_input)
        break